import os
from glob import glob
from os.path import join, exists

from torch import nn
from torch.utils.tensorboard import SummaryWriter

from config import conf

MODEL_DIR = conf.get('dir', 'model_save_dir')
LOG_DIR = conf.get('dir', 'log_dir')


class BaseModel(nn.Module):
    def __init__(self, model_name, version_id=None, continue_train=False):
        """
        模型的基类
        :param model_name: 要保存的模型的名称
        :param version_id:仅当要加载已经训练好的某个模型时，才指定version_id；正常训练的version_id会自动生成
        """
        super().__init__()
        model_save_dir, log_dir, version_id, model_root = self.prepare_model_dir(model_name, version_id)
        self.model_root = model_root
        self.log_dir = log_dir
        self.version_id = version_id
        self.model_save_dir = model_save_dir
        board_dir = join(log_dir, 'board')
        self.writer = SummaryWriter(board_dir)
        self.model_name = model_name
        print(model_name, 'version%d' % version_id)
        print('tensorboard:%s' % (board_dir))

    def prepare_model_dir(self, model_name, version_id):
        """
        创建模型的保存路径并返回模型的版本
        :param model_name:
        :param version_id:
        :return:
        """
        model_root = join(MODEL_DIR, model_name)
        if not exists(model_root):
            os.makedirs(model_root)

        if version_id is not None:
            model_save_dir = join(model_root, 'version%s' % version_id)
            if not exists(model_save_dir):
                raise Exception('model %s does not exist' % (version_id))
        else:
            version_id = len(glob(join(model_root, 'version*'))) + 1
            model_save_dir = join(model_root, 'version%s' % version_id)
            os.mkdir(model_save_dir)

        # 准备log dir
        log_dir = join(model_root, 'board', str(version_id))
        if not exists(log_dir):
            os.makedirs(log_dir)
        if not exists(join(model_save_dir, 'data')):
            os.mkdir(join(model_save_dir, 'data'))
        return model_save_dir, log_dir, version_id, model_root
